"""并行生成日度 / 月度 时间序列动图（不保存中间静态图）

用法（示例）：
  ./.venv/bin/python scripts/generate_time_series_gif.py \
      --parquet data/云南乡镇月度灯光亮度.parquet \
      --out figures/月度灯光_test.gif --limit 6

特性：
- 自动尝试检测时间列和数值列（若未指定）。
- 并行渲染每帧（ProcessPoolExecutor），只在内存中生成 PNG 并合成 GIF，不保存中间静态图。
- 支持 --limit 限制帧数，便于 smoke test。
"""
from pathlib import Path
import sys

# Ensure project root is on sys.path so `import src` works when running from scripts/
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))
import argparse
import json
import io
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import List, Optional, Tuple
import multiprocessing
import math

import pandas as pd
import geopandas as gpd
from PIL import Image as PILImage

from src.aux_utils import merge_geo_population
from utils.plot_style import set_plot_style, set_chinese_font
import matplotlib as mpl

# Import colorbar formatters from existing module (they are module-level helpers)
from src.gif_generator import _format_cb_ticks_identity, _format_cb_ticks_percent, _format_cb_ticks_percent_identity


def _ts_worker(task):
    """Worker run in subprocess: draws one frame and returns (label, png_bytes).

    task: dict with keys: label(str), features(list), params(dict)
    """
    label = task['label']
    features = task['features']
    params = task['params']
    mpl.use('Agg')
    import matplotlib.pyplot as plt
    from shapely.geometry import shape
    from matplotlib.ticker import FuncFormatter
    # local imports for style
    from utils.plot_style import set_plot_style, set_chinese_font

    props = [feat['properties'] for feat in features]
    geoms = [shape(feat['geometry']) for feat in features]
    import pandas as pd
    import geopandas as gpd
    df = pd.DataFrame(props)
    gdf = gpd.GeoDataFrame(df, geometry=geoms, crs=params.get('crs'))

    set_plot_style(params.get('plot_cfg', {}))
    font_family = params.get('font_family')
    if font_family:
        set_chinese_font(font_family, quiet=True)

    figsize = params.get('figsize', (10, 12))
    dpi = params.get('dpi', 150)
    column = params.get('column')
    cmap = params.get('cmap', 'Reds')
    vmin = params.get('vmin')
    vmax = params.get('vmax')
    colorbar_bbox = params.get('colorbar_bbox', [0.15, 0.04, 0.7, 0.03])

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    if column in gdf.columns:
        col = column
    else:
        col = 'mean' if 'mean' in gdf.columns else '总人口'

    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    gdf.plot(column=col, cmap=cmap, linewidth=params.get('linewidth', 0.2), edgecolor=params.get('edgecolor', '#444444'), ax=ax, norm=norm)
    ax.set_axis_off()

    mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    mappable._A = []
    cax = fig.add_axes(params.get('colorbar_bbox', colorbar_bbox))
    cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

    tick_percent = params.get('colorbar_tick_percent', False)
    tick_percent_mult = params.get('colorbar_tick_percent_multiply', True)
    if tick_percent:
        if tick_percent_mult:
            cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_percent))
        else:
            cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_percent_identity))
    else:
        cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_identity))
    cb.ax.tick_params(labelsize=9)

    label = params.get('colorbar_label', '')
    if tick_percent and ('%' not in label and '（' not in label and '(' not in label):
        label = f"{label}（%）"
    cb.set_label(label, fontsize=params.get('colorbar_label_size', 11))

    year_label_kwargs = dict(color=params.get('frame_label_color', 'white'), fontsize=params.get('frame_label_size', 18), weight='bold', va='bottom', ha='right')
    if font_family:
        year_label_kwargs['fontfamily'] = font_family
    ax.text(0.97, 0.05, str(label if params.get('use_label_as_cb', False) else task['label']), transform=ax.transAxes, bbox=dict(facecolor='black', alpha=0.35, pad=4), **year_label_kwargs)

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buf.seek(0)
    data = buf.getvalue()
    buf.close()
    return task['label'], data


def detect_columns(df: pd.DataFrame) -> Tuple[Optional[str], Optional[str]]:
    # time candidates (Chinese and English)
    time_cands = ['日期', '日', 'time', '时间', 'timestamp', '月份', 'month', '年月', '年月日']
    val_cands = ['mean', 'value', '亮度', 'avg', 'radiance', '值', '亮度均值']
    time_col = None
    val_col = None
    for c in df.columns:
        if c in time_cands:
            time_col = c
            break
    for c in df.columns:
        if c in val_cands:
            val_col = c
            break
    # fallback: first column that looks like datetime
    if time_col is None:
        for c in df.columns:
            if pd.api.types.is_datetime64_any_dtype(df[c]) or pd.api.types.is_object_dtype(df[c]):
                # try parseable
                try:
                    pd.to_datetime(df[c].iloc[:10])
                    time_col = c
                    break
                except Exception:
                    continue
    return time_col, val_col


def main(argv: Optional[List[str]] = None):
    p = argparse.ArgumentParser(description='生成月度/日度时间序列 GIF（并行，无中间静态图）')
    p.add_argument('--parquet', required=True)
    p.add_argument('--geojson', default='data/云南乡镇.json')
    p.add_argument('--time-col', default=None)
    p.add_argument('--value-col', default=None)
    p.add_argument('--merge-on', default=None, help='用于合并的 code 字段（如 town_code），默认自动检测')
    p.add_argument('--out', required=True)
    p.add_argument('--limit', type=int, default=0, help='仅用于测试：限制渲染的帧数（取最前面的 n 个时间点）')
    p.add_argument('--n-workers', type=int, default=0)
    p.add_argument('--duration', type=float, default=0.8)
    p.add_argument('--dpi', type=int, default=150)
    p.add_argument('--figsize', type=float, nargs=2, default=(10, 12))
    p.add_argument('--cmap', default='Reds')
    p.add_argument('--colorbar-label', default='值')
    p.add_argument('--colorbar-percent', action='store_true')
    args = p.parse_args(argv)

    df = pd.read_parquet(args.parquet)
    time_col = args.time_col
    val_col = args.value_col
    if time_col is None or val_col is None:
        det_time, det_val = detect_columns(df)
        time_col = time_col or det_time
        val_col = val_col or det_val
    if time_col is None:
        raise RuntimeError('无法检测到时间列，请使用 --time-col 指定')
    if val_col is None:
        raise RuntimeError('无法检测到数值列，请使用 --value-col 指定')

    # parse time column to datetime for sorting
    try:
        df[time_col] = pd.to_datetime(df[time_col])
    except Exception:
        # leave as is
        pass

    # load geo
    geo = gpd.read_file(args.geojson)

    # determine merge key
    merge_on = args.merge_on
    if merge_on is None:
        if '乡镇代码' in df.columns or 'town_code' in df.columns or 'tid' in df.columns:
            merge_on = 'town_code'
        else:
            merge_on = None

    # build times list
    times = sorted(df[time_col].dropna().unique())
    if args.limit and args.limit > 0:
        times = times[: args.limit]

    # compute global vmin/vmax across selected times (use val_col)
    all_vals = []
    for t in times:
        sub = df[df[time_col] == t]
        if val_col in sub.columns:
            vals = pd.to_numeric(sub[val_col], errors='coerce').dropna().values
            if vals.size > 0:
                all_vals.append(vals)
    if all_vals:
        import numpy as np
        stacked = np.concatenate(all_vals)
        vmin = float(stacked.min())
        vmax = float(np.percentile(stacked, 95))
        if vmin == vmax:
            vmax = vmin + 1e-6
    else:
        vmin, vmax = 0.0, 1.0

    # prepare tasks: for each time, merge with geo and serialize features
    tasks = []
    for t in times:
        sub = df[df[time_col] == t]
        # ensure merge key exists in sub
        merged = merge_geo_population(geo, sub, merge_on_code='town_code' if 'town_code' in sub.columns else merge_on, name_suffixes=None, normalize_name_fn=None)
        try:
            features = json.loads(merged.to_json())['features']
        except Exception:
            features = None
            # fallback: use geometry __geo_interface__
            try:
                features = merged.__geo_interface__['features']
            except Exception:
                raise
        params = {
            'figsize': tuple(args.figsize),
            'dpi': args.dpi,
            'column': val_col,
            'cmap': args.cmap,
            'vmin': vmin,
            'vmax': vmax,
            'colorbar_bbox': (0.15, 0.04, 0.7, 0.03),
            'plot_cfg': {},
            'font_family': None,
            'colorbar_label': args.colorbar_label,
            'colorbar_tick_percent': args.colorbar_percent,
            'colorbar_tick_percent_multiply': True,
            'colorbar_label_size': 11,
            'frame_label_color': 'white',
            'frame_label_size': 18,
            'linewidth': 0.2,
            'edgecolor': '#444444',
        }
        tasks.append({'label': str(t), 'features': features, 'params': params})

    # run parallel worker
    n_workers = args.n_workers if args.n_workers and args.n_workers > 0 else (multiprocessing.cpu_count() or 1)
    results = []
    with ProcessPoolExecutor(max_workers=n_workers) as ex:
        futures = {ex.submit(_ts_worker, task): task['label'] for task in tasks}
        for fut in as_completed(futures):
            try:
                label, data = fut.result()
                results.append((label, data))
            except Exception:
                import traceback
                print('Worker exception:', traceback.format_exc())

    # sort by label order in tasks
    label_order = {task['label']: i for i, task in enumerate(tasks)}
    results_sorted = sorted(results, key=lambda x: label_order.get(x[0], 0))

    frames = []
    for lbl, data in results_sorted:
        buf = io.BytesIO(data)
        pil_img = PILImage.open(buf).convert('RGBA')
        frames.append(pil_img.copy())
        buf.close()

    if not frames:
        raise RuntimeError('没有渲染到任何帧')

    outp = Path(args.out).resolve()
    outp.parent.mkdir(parents=True, exist_ok=True)
    duration_ms = int(args.duration * 1000)
    frames[0].save(outp, save_all=True, append_images=frames[1:], duration=duration_ms, loop=0, optimize=False, disposal=2)
    print(f'已保存 GIF 到 {outp} (frames={len(frames)})')


if __name__ == '__main__':
    main()
